#!/bin/bash

# 函数：显示用法
usage() {
    echo -e "Usage: $0 [options...]"
    echo -e " --path <string>                                [Required] The folder to be monitored"
    echo -e " --api-url <string>                             [Required] Dify's API interface base URL"
    echo -e " --api-key <string>                             [Required] Dify's API key"
    echo -e " --dataset-id <string>                          [Required] Dify's knowledge database_id"
    echo -e ""
    echo -e " --indexing-technique <string>                  [option] Index method, use high_quality (Default) / economy"
    echo -e " --doc-form <string>                            [option] Format of Index Content, use text_model (Default) / hierarchical_model / qa_model"
    echo -e " --doc-language <string>                        [option] If doc-form use qa_model, it can use English / Chinese"
    echo -e " --process-rule-mode <string>                   [option] Cleaning and segmentation mode, use automatic (Default) / custom"
    echo -e " --pre-processing-rules <string>                [option] If process-rule-mode use custom, it can use remove_extra_spaces / remove_urls_emails"
    echo -e " --segmentation-separator <string>              [option] Custom segment identifier"
    echo -e " --segmentation-max-tokens <number>           	[option] Max segment max size"
    echo -e " --parent-mode <string>                         [option] Recall mode of parent segmentation, use full-doc / paragraph"
    echo -e " --subchunk-segmentation-separator <string>     [option] Child segment identifier"
    echo -e " --subchunk-segmentation-max-tokens <number>    [option] Child segment max size"
    echo -e " --log-file <string>                            [option] Logfile, save to the same directory by default"
    exit 1
}

# 验证输入参数
if [[ $# -eq 0 ]]; then
    # 获取 max_queued_events 的值
    max_queued_events=$(sysctl -n fs.inotify.max_queued_events)


    # 输出中间部分，包含内容
    echo -e "\e[33mNotice: Current max_queued_events is $max_queued_events, Maximum support for submitting $max_queued_events file changes at once.\nYou can use 'sudo sysctl -w fs.inotify.max_queued_events=32768' to change it temporarily. \nor use echo 'fs.inotify.max_queued_events=32768' | sudo tee -a /etc/sysctl.conf to change permanently.\e[0m\n\ntry '$0 --help' for more information"
    exit 1
fi

# 解析参数
for arg in "$@"; do
    case $arg in
        --help)
            usage
            ;;
        --path=*)
            PATH_TO_MONITOR="${arg#*=}"
            ;;
        --api-url=*)
            API_URL="${arg#*=}"
            ;;
        --api-key=*)
            API_KEY="${arg#*=}"
            ;;
        --dataset-id=*)
            DATASET_ID="${arg#*=}"
            ;;
        --indexing-technique=*)
            INDEXING_TECHNIQUE="${arg#*=}"
            ;;
        --doc-form=*)
            DOC_FORM="${arg#*=}"
            ;;
        --doc-language=*)
            DOC_LANGUAGE="${arg#*=}"
            ;;
        --process-rule-mode=*)
            PROCESS_RULE_MODE="${arg#*=}"
            ;;
        --pre-processing-rules=*)
            PRE_PROCESSING_RULES="${arg#*=}"
            ;;
        --segmentation-separator=*)
            SEGMENTATION_SEPARATOR="${arg#*=}"
            ;;
        --segmentation-max-tokens=*)
            SEGMENTATION_MAX_TOKENS="${arg#*=}"
            ;;
        --parent-mode=*)
            PARENT_MODE="${arg#*=}"
            ;;
        --subchunk-segmentation-separator=*)
            SUBCHUNK_SEGMENTATION_SEPARATOR="${arg#*=}"
            ;;
        --subchunk-segmentation-max-tokens=*)
            SUBCHUNK_SEGMENTATION_MAX_TOKENS="${arg#*=}"
            ;;
        --log-file=*)
            LOG_FILE="${arg#*=}"
            ;;
        *)
            echo "Unknown option: $arg"
            exit 1
            ;;
    esac
done

# 验证必需参数
if [[ -z "$PATH_TO_MONITOR" ]]; then
    echo "Error: --path is required."
    exit 1
fi

if [[ -z "$API_URL" ]]; then
    echo "Error: --api-url is required."
    exit 1
fi

if [[ -z "$API_KEY" ]]; then
    echo "Error: --api-key is required."
    exit 1
fi

if [[ -z "$DATASET_ID" ]]; then
    echo "Error: --dataset-id is required."
    exit 1
fi

# 设置默认值
INDEXING_TECHNIQUE=${INDEXING_TECHNIQUE:-high_quality}
DOC_FORM=${DOC_FORM:-text_model}
PROCESS_RULE_MODE=${PROCESS_RULE_MODE:-automatic}
DOC_LANGUAGE=${DOC_LANGUAGE:-Chinese}
SEGMENTATION_MAX_TOKENS=${SEGMENTATION_MAX_TOKENS:-1000}
PARENT_MODE=${PARENT_MODE:-paragraph}
SUBCHUNK_SEGMENTATION_MAX_TOKENS=${SUBCHUNK_SEGMENTATION_MAX_TOKENS:-500}
LOG_FILE=${LOG_FILE:-dify_dataset_syn.log}

# 定义允许的取值
allowed_indexing=("high_quality" "economy")
allowed_doc_form=("text_model" "hierarchical_model" "qa_model")
allowed_process_rule=("automatic" "custom")
allowed_doc_language=("English" "Chinese")
allowed_parent_mode=("full-doc" "paragraph")

# 校验函数：判断给定的值是否在数组中
function check_value() {
    local value="$1"
    shift
    local allowed_values=("$@")
    for v in "${allowed_values[@]}"; do
        if [[ "$v" == "$value" ]]; then
            return 0
        fi
    done
    return 1
}


# 检查 INDEXING_TECHNIQUE 输入
if ! check_value "$INDEXING_TECHNIQUE" "${allowed_indexing[@]}"; then
    echo "Error: The value '$INDEXING_TECHNIQUE' for INDEXING_TECHNIQUE is invalid. Allowed values are: [${allowed_indexing[*]}]."
    exit 1
fi

# 检查 DOC_FORM 输入
if ! check_value "$DOC_FORM" "${allowed_doc_form[@]}"; then
    echo "Error: The value '$DOC_FORM' for DOC_FORM is invalid. Allowed values are: [${allowed_doc_form[*]}]."
    exit 1
fi

# 检查 PROCESS_RULE_MODE 输入
if ! check_value "$PROCESS_RULE_MODE" "${allowed_process_rule[@]}"; then
    echo "Error: The value '$PROCESS_RULE_MODE' for PROCESS_RULE_MODE is invalid. Allowed values are: [${allowed_process_rule[*]}]."
    exit 1
fi

# 检查 DOC_LANGUAGE 输入
if ! check_value "$DOC_LANGUAGE" "${allowed_doc_language[@]}"; then
    echo "Error: The value '$DOC_LANGUAGE' for DOC_LANGUAGE is invalid. Allowed values are: [${allowed_doc_language[*]}]."
    exit 1
fi

# 检查 PARENT_MODE 输入
if ! check_value "$PARENT_MODE" "${allowed_parent_mode[@]}"; then
    echo "Error: The value '$PARENT_MODE' for PARENT_MODE is invalid. Allowed values are: [${allowed_parent_mode[*]}]."
   exit 1
fi

# 验证路径是否存在
if [[ ! -d "$PATH_TO_MONITOR" ]]; then
    echo "Error: The specified path '$PATH_TO_MONITOR' is not a directory."
    exit 1
fi

# 检查是否安装inotifywait
if ! command -v inotifywait &> /dev/null; then
    echo "Error: inotify is not installed. use apt install inotify-tools to fix it."
    exit 1
fi

# 检查是否安装jq
if ! command -v jq &> /dev/null; then
    echo "Error: jq is not installed. use apt install jq to fix it."
    exit 1
fi

# 函数：query_file_by_name
# 参数：
#   $1 - 文件名称
#   $2 - dataset_id
#   $3 - api_key
# 功能：通过调用 API 查询文件信息，返回文件 id
# 使用：$(query_file_by_name "$file_name")
query_file_by_name() {
    local file_name="$1"
    # 构造请求 URL
    local url="$API_URL/datasets/$DATASET_ID/documents?keyword=$file_name"

    # 调用 API 并获取返回的 JSON 数据
    local response
    response=$(curl --silent --location --request GET "$url" \
      --header "Authorization: Bearer $API_KEY")

    # 判断 data 数组的长度是否为 0
    if [ "$(echo "$response" | jq '.data | length')" -eq 0 ]; then
        echo "Error: can not find document ${file_name}" >&2
        return 1
    fi

    # 通过 jq 筛选出数组中 name 字段和传入 file_name 匹配的项，获取 id
    local file_id
    file_id=$(echo "$response" | jq -r --arg name "$file_name" '.data[] | select(.name == $name) | .id')

    if [ -z "$file_id" ]; then
        echo "Error: can not find document '${file_name}' ID" >&2
        return 1
    fi

    echo "$file_id"
    return 0
}

echo "Monitoring changes in $PATH_TO_MONITOR..."

handle_event() {
    case $event in
         # 新建文件/复制文件/从不同文件系统移动文件事件
        MODIFY|MOVED_TO)
            echo "[$(date +"%Y-%m-%d %H:%M:%S")] UPLOAD FILE $file: Get notified." | tee -a $LOG_FILE
            # 通过 jq 构造 JSON 数据  
            payload=$(jq -n \
                --arg indexing_technique "$INDEXING_TECHNIQUE" \
                --arg doc_form "$DOC_FORM" \
                --arg doc_language "$DOC_LANGUAGE" \
                --arg process_rule_mode "$PROCESS_RULE_MODE" \
                --arg pre_processing_rules "$PRE_PROCESSING_RULES" \
                --arg segmentation_separator "$SEGMENTATION_SEPARATOR" \
                --arg segmentation_max_tokens $SEGMENTATION_MAX_TOKENS \
                --arg parent_mode "$PARENT_MODE" \
                --arg subchunk_segmentation_separator "$SUBCHUNK_SEGMENTATION_SEPARATOR" \
                --arg subchunk_segmentation_max_tokens $SUBCHUNK_SEGMENTATION_MAX_TOKENS \
                '{
                    indexing_technique: $indexing_technique,
                    doc_form: $doc_form,
                    doc_language: $doc_language,
                    process_rule: {
                        mode: $process_rule_mode,
                        rules: {
                            "pre_processing_rules": ($pre_processing_rules | split(",") | map({id: ., enabled: true})),
                            segmentation: {
                                separator: $segmentation_separator, 
                                max_tokens: ($segmentation_max_tokens | tonumber)
                            },
                            parent_mode: $parent_mode,
                            subchunk_segmentation: {
                                separator: $subchunk_segmentation_separator,
                                max_tokens: ($subchunk_segmentation_max_tokens | tonumber)
                            }
                        }
                    }
                }')
            # 发送 RESET 请求
            response=$(curl --location --request POST "$API_URL/datasets/$DATASET_ID/document/create-by-file" \
                --header "Authorization: Bearer $API_KEY" \
                --form "data=\"$(printf %s "$payload" | sed 's/"/\\"/g')\"" \
                --form "file=@\"$PATH_TO_MONITOR/$file\"")

            # 验证curl返回	
            code=$(echo "$response" | jq -r '.code')
            message=$(echo "$response" | jq -r '.message')

            if [[ "$code" != "null" && "$code" != "0" ]]; then
                echo "[$(date +"%Y-%m-%d %H:%M:%S")] UPLOAD FILE $file: $message" | tee -a $LOG_FILE
            else
                echo "[$(date +"%Y-%m-%d %H:%M:%S")] UPLOAD FILE $file: Syn successful" | tee -a $LOG_FILE
            fi
            echo "[$(date +"%Y-%m-%d %H:%M:%S")] \`\`\`\`\`\`"
            ;;

        # 删除/移出文件事件
        DELETE|MOVED_FROM)
            echo "delete $file"
            echo "[$(date +"%Y-%m-%d %H:%M:%S")] DELETE FILE $file: Get notified." | tee -a $LOG_FILE
            document_id=$(query_file_by_name "$file")
            response=$(curl --location --request DELETE "$API_URL/datasets/$DATASET_ID/documents/$document_id" \
                --header "Authorization: Bearer $API_KEY")
            
            # 验证curl返回	
            code=$(echo "$response" | jq -r '.code')
            message=$(echo "$response" | jq -r '.message')

            if [[ "$code" != "null" && "$code" != "0" ]]; then
                echo "[$(date +"%Y-%m-%d %H:%M:%S")] DELETE FILE $file: $file$message" | tee -a $LOG_FILE
            else
                echo "[$(date +"%Y-%m-%d %H:%M:%S")] DELETE FILE $file: Syn successful" | tee -a $LOG_FILE
            fi
            echo "[$(date +"%Y-%m-%d %H:%M:%S")] \`\`\`\`\`\`"
            ;;
    esac
}

# 使用inotifywait监控文件变动
inotifywait -m -r -e delete,modify,moved_to,moved_from "$PATH_TO_MONITOR" | while read -r directory event file; do
    handle_event <<< "$event"
done